#!/usr/bin/env python3
# _*_ coding:utf-8 _*_

import math
import torch
import torch.nn as nn
import torch.nn.init as init
import scipy
import numpy as np
import matplotlib.pyplot as plt
from options import Options

PI = np.pi
class Activation(nn.Module):
    def __init__(self, act_name):
        super().__init__()
        self.act_name = act_name

    def forward(self, x):
        if self.act_name == 'tanh':
            return torch.tanh(x)
        elif self.act_name == 'cos':
            return torch.cos(np.pi * x)
        elif self.act_name == 'sin':
            return torch.sin(np.pi * x)


class ResBlock(nn.Module):
    def __init__(self, hidden_dim, act_name):
        super().__init__()
        self.act_name = act_name
        self.hidden_dim = hidden_dim
        self.main = nn.Sequential(nn.Linear(self.hidden_dim, self.hidden_dim),
                                  Activation(self.act_name))

    def forward(self, X):
        out = self.main(X)
        return out + X


class GaussSplattingPde1D(nn.Module):
    def __init__(self, args):
        super().__init__()
        self.dim = 1
        self.args = args
        self.weights = nn.Parameter(torch.ones(1, self.args.n_gauss) - 2 * torch.rand(1, self.args.n_gauss))   # torch.ones(1, self.args.n_gauss)
        self.mean = nn.Parameter(torch.rand(1, self.args.n_gauss) * 2 * PI)
        self.std = nn.Parameter(torch.rand(1, self.args.n_gauss))        

    def _name(self):
        return f'GaussSplattingPde1D_model_gauss_{self.args.n_gauss}_n_{self.args.n_pde}_lr_{self.args.lr}_case_{self.args.case}'

    def forward(self, X):
        gauss_f = torch.exp(- 0.5 * (X - self.mean) ** 2 / self.std ** 2)
        return torch.sum(self.weights * gauss_f, dim=1, keepdim=True)


if __name__ == '__main__':
    args = Options().parse()
    args.n_gauss = 4
    model = GaussSplattingPde1D(args)
    # for param in model.parameters():
    #     print(type(param.data), param.size())
    # for k, v in model.state_dict().items():
    #     print(k, v.size())

    X = torch.rand((3,1))
    print("X", X)
    print("weights", model.weights)
    print("points", model.points)
    print("mean", model.mean)
    print("std", model.std)

    print(X.shape, model.mean.shape)
    a = X-model.mean
    b = torch.exp(- 0.5 * (X - model.mean) ** 2 / model.std ** 2)

    print(a)
    print(b)
    print(model.weights * b)
    print(model.forward(X))
